# 功能：加载原始市场票价数据，并记录元信息与数据质量报告

import pandas as pd
import os
from datetime import datetime


def load_data(file_path, sample_size=None, output_path=None):
    """
    加载原始市场票价数据，并生成数据质量报告
    :param file_path: 原始CSV文件路径
    :param sample_size: 是否采样（如取前N条），默认不采样
    :param output_path: 输出Parquet文件路径（若为None则不保存）
    :return: DataFrame格式数据
    """
    print("【开始】加载原始数据...")

    # 记录采集时间与参数
    start_time = datetime.now()
    print(f" 开始时间：{start_time}")

    # 加载数据
    df = pd.read_csv(file_path)
    original_shape = df.shape
    print(f" 原始数据加载完成：共 {original_shape[0]} 行，{original_shape[1]} 列")

    # 采样逻辑（可选）
    if sample_size:
        df = df.head(sample_size).copy()
        print(f" 已采样前 {sample_size} 条数据")

    # 数据质量初步分析
    missing_summary = df.isnull().sum()
    missing_rate = (missing_summary / df.shape[0]) * 100
    quality_report = pd.DataFrame({
        'Missing Count': missing_summary,
        'Missing Rate (%)': missing_rate,
        'Unique Count': df.nunique(),
        'Data Type': df.dtypes
    })
    print("\n 数据质量报告：")
    print(quality_report)

    # 元信息记录
    metadata = {
        "source_file": file_path,
        "sample_size": sample_size if sample_size else original_shape[0],
        "columns_count": original_shape[1],
        "loaded_time": start_time.strftime("%Y-%m-%d %H:%M:%S"),
        "output_file": output_path if output_path else "Not saved"
    }
    print("\n 元信息记录：")
    for key, value in metadata.items():
        print(f" - {key}: {value}")

    # 保存为 Parquet（可选）
    if output_path:
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        df.to_parquet(output_path)
        print(f" 数据已保存至：{output_path}")

    print("【完成】数据加载流程结束")
    return df


if __name__ == "__main__":
    # 示例运行路径
    input_path = "../data/processed/MarketFarePredictionData.csv"
    output_path = "../data/raw/raw_data.parquet"

    # 调用函数加载并保存数据
    raw_data = load_data(input_path, sample_size=100_000, output_path=output_path)

    # 打印前几行查看数据
    print("\n 数据预览：")
    print(raw_data.head())
